{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "from transformers import BertTokenizer, BertConfig\n",
    "from transformers import BertModel\n",
    "import nltk\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import re\n",
    "import nltk\n",
    "from sklearn import model_selection \n",
    "from sklearn.linear_model import LogisticRegression\n",
    "import pickle\n",
    "import string\n",
    "from transformers import BertTokenizer, BertConfig\n",
    "from transformers import BertModel\n",
    "import torch\n",
    "from nltk.corpus import stopwords\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_token_positions(input_text,token_list,start_from_pos=0):\n",
    "    pos_list = []                    \n",
    "    name_to_match = input_text.lower().replace(\" \",\"\")\n",
    "    remaining_name = input_text.lower().replace(\" \",\"\")\n",
    "    name = \"\"\n",
    "    count = start_from_pos\n",
    "    for entry in token_list[start_from_pos:]:\n",
    "        if(remaining_name.startswith(entry.strip(\"##\").lower())):\n",
    "            pos_list.append(count)\n",
    "            remaining_name = remaining_name[len(entry.strip(\"##\").lower()):]\n",
    "            name = name + entry.strip(\"##\").lower()\n",
    "            if(name == name_to_match):\n",
    "                break\n",
    "        else:\n",
    "            pos_list = []\n",
    "            name = \"\"\n",
    "            remaining_name = name_to_match\n",
    "            if(remaining_name.startswith(entry.strip(\"##\").lower())):                                 \n",
    "                pos_list.append(count)                                                                \n",
    "                remaining_name = remaining_name[len(entry.strip(\"##\").lower()):]\n",
    "                name = name + entry.strip(\"##\").lower()    \n",
    "                if(name == name_to_match):                                   \n",
    "                    break\n",
    "\n",
    "        count = count + 1\n",
    "\n",
    "    return pos_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_string(string_input, padding_length):\n",
    "    #string_input = re.sub(r\"(\\w)([\"+ string.punctuation +\"])\", r\"\\1 \\2 \", string_input)\n",
    "    #string_input = string_input.replace(\".\",\" . \")\n",
    "\n",
    "    tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')\n",
    "    sentences = tokenizer.tokenize(string_input)\n",
    "    \n",
    "    config = BertConfig.from_pretrained('C:/Users/itsma/Documents/BERT_models/NCBI_BERT_pubmed_mimic_uncased_L-12_H-768_A-12')\n",
    "    config.output_hidden_states = True\n",
    "\n",
    "    bert_tokenizer = BertTokenizer.from_pretrained('C:/Users/itsma/Documents/BERT_models/NCBI_BERT_pubmed_mimic_uncased_L-12_H-768_A-12')\n",
    "    bert_model = BertModel.from_pretrained(\"C:/Users/itsma/Documents/BERT_models/NCBI_BERT_pubmed_mimic_uncased_L-12_H-768_A-12\",config=config)\n",
    "    #bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "    #bert_model = BertModel.from_pretrained('bert-base-uncased')\n",
    "    #stop_words = set(stopwords.words('english')) \n",
    "    \n",
    "    positions_covered = 0\n",
    "    word_list = []\n",
    "    \n",
    "    for index in range(len(sentences)):\n",
    "        new_dict_sentence = {}\n",
    "        sentence = sentences[index]\n",
    "        new_dict_sentence[\"sentence\"] = sentence\n",
    "        new_dict_sentence[\"padding_length\"] = padding_length\n",
    "        start_index_bert = max(0,index-padding_length)\n",
    "        end_index_bert = min(len(sentences),index+padding_length)\n",
    "\n",
    "        bert_input = ' '.join(sentences[start_index_bert:(end_index_bert+1)])\n",
    "\n",
    "        input_ids = torch.tensor(bert_tokenizer.encode(bert_input)).unsqueeze(0)  \n",
    "        outputs = bert_model(input_ids)\n",
    "        bert_vector = outputs[2][0][0].data.numpy()\n",
    "        bert_tokens = bert_tokenizer.tokenize(bert_input)\n",
    "\n",
    "        '''\n",
    "        bert_vec_list = []\n",
    "        for entry in range(0,len(outputs[0][0])):\n",
    "            bert_vec_list.append(outputs[0][0][entry].data.numpy())\n",
    "\n",
    "        bert_vector_sentence = np.mean(bert_vec_list,axis=0)\n",
    "        '''\n",
    "        start_pos = 0\n",
    "        prior_pos = get_token_positions(' '.join(sentences[start_index_bert:index]),bert_tokens)\n",
    "        \n",
    "        if(len(prior_pos)>0):\n",
    "            start_pos = max(prior_pos)\n",
    "            \n",
    "        tokens = nltk.word_tokenize(sentence)\n",
    "        pos_tokens = nltk.pos_tag(tokens)\n",
    "\n",
    "        sentence_covered = ''\n",
    "        \n",
    "        for token in pos_tokens:\n",
    "            new_dict = {}\n",
    "            current_word = token[0]\n",
    "            new_dict[\"word\"] = current_word\n",
    "            token_position = string_input.find(current_word, positions_covered)\n",
    "            spaces_between = string_input[positions_covered:token_position] \n",
    "            sentence_covered = sentence_covered + spaces_between + current_word\n",
    "            positions_covered = token_position + len(current_word)\n",
    "            new_dict[\"begin_pos\"] = token_position\n",
    "            new_dict[\"end_pos\"] = positions_covered\n",
    "\n",
    "            bert_token_positions = get_token_positions(current_word,bert_tokens,start_pos)\n",
    "            vec_list = []\n",
    "            if(len(bert_token_positions)==0):\n",
    "                continue\n",
    "            start_pos = bert_token_positions[-1] + 1\n",
    "\n",
    "            for entry in bert_token_positions:\n",
    "                vec_list.append(bert_vector[entry])\n",
    "            \n",
    "            bert_vector_word = np.mean(vec_list,axis=0)\n",
    "            new_dict[\"keyword_vector\"] = bert_vector_word\n",
    "\n",
    "            word_list.append(new_dict)\n",
    "    \n",
    "    return word_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_available_in_pos_list(dict_list, begin_pos, end_pos):\n",
    "    for entry in dict_list:\n",
    "        if(entry[\"start\"]<=begin_pos and entry[\"end\"]>=end_pos):\n",
    "            return True\n",
    "    \n",
    "    return False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discharge_note():\n",
    "    def __init__(self, root):\n",
    "        self.xml_root = root\n",
    "        \n",
    "    def process_note(self):\n",
    "        root = self.xml_root\n",
    "        text_section = root.find('TEXT')\n",
    "        text = text_section.text\n",
    "        \n",
    "        self.processed_text = process_string(text,1)\n",
    "\n",
    "        tag_section = root.find('TAGS')\n",
    "        event_list = []\n",
    "        timex_list = []\n",
    "        tlink_list = []\n",
    "        sectime_list = []\n",
    "        for child in tag_section:\n",
    "            if(child.tag=='EVENT'):\n",
    "                event_list.append(child.attrib)\n",
    "            elif(child.tag=='TIMEX3'):\n",
    "                timex_list.append(child.attrib)\n",
    "            elif(child.tag=='TLINK'):\n",
    "                tlink_list.append(child.attrib)\n",
    "            elif(child.tag=='SECTIME'):\n",
    "                sectime_list.append(child.attrib)\n",
    "\n",
    "        for sub in event_list:\n",
    "            sub[\"start\"] = int(sub[\"start\"])\n",
    "            sub[\"end\"] = int(sub[\"end\"])\n",
    "\n",
    "        for sub in timex_list:\n",
    "            sub[\"start\"] = int(sub[\"start\"])\n",
    "            sub[\"end\"] = int(sub[\"end\"])\n",
    "\n",
    "        for sub in sectime_list:\n",
    "            sub[\"start\"] = int(sub[\"start\"])\n",
    "            sub[\"end\"] = int(sub[\"end\"])\n",
    "\n",
    "\n",
    "        event_list = sorted(event_list, key = lambda i: i['start'])\n",
    "        timex_list = sorted(timex_list, key = lambda i: i['start'])\n",
    "        sectime_list = sorted(sectime_list, key = lambda i: i['start'])\n",
    "    \n",
    "        for entry in self.processed_text:\n",
    "            event_entry_available = is_available_in_pos_list(event_list,entry[\"begin_pos\"],entry[\"end_pos\"])\n",
    "            timex_entry_available = is_available_in_pos_list(timex_list,entry[\"begin_pos\"],entry[\"end_pos\"])\n",
    "            sectime_entry_available = is_available_in_pos_list(sectime_list,entry[\"begin_pos\"],entry[\"end_pos\"])\n",
    "\n",
    "            entry.update({\"event_flag\":1 if event_entry_available else 0})\n",
    "            entry.update({\"timex_flag\":1 if timex_entry_available else 0})\n",
    "            entry.update({\"sectime_flag\":1 if sectime_entry_available else 0})\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_list = []\n",
    "count = 0\n",
    "for file in os.listdir(\"C:/Users/itsma/Documents/Capstone project/2012-07-15.original-annotation.release/\"):\n",
    "    if file.endswith(\".xml\"):\n",
    "        try:\n",
    "            file_name = os.path.join(\"C:/Users/itsma/Documents/Capstone project/2012-07-15.original-annotation.release/\", file)\n",
    "            tree = ET.parse(file_name)\n",
    "            root = tree.getroot()\n",
    "            discharge_note = Discharge_note(root)\n",
    "            discharge_note.process_note()\n",
    "            word_list.extend(discharge_note.processed_text)\n",
    "            count = count + 1\n",
    "            print(count)\n",
    "            \n",
    "        except:\n",
    "            count = count + 1\n",
    "            print(\"Error in processing file:\",file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(word_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_list_df = pd.DataFrame(word_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_data = word_list_df[word_list_df[\"sectime_flag\"]==1]\n",
    "neg_data = word_list_df[word_list_df[\"sectime_flag\"]==0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_data_volume = len(pos_data)\n",
    "neg_data = neg_data.sample(pos_data_volume)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(pos_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_data = pd.concat([pos_data,neg_data])\n",
    "X = np.vstack(list(combined_data[\"keyword_vector\"]))                                \n",
    "y = combined_data.sectime_flag  \n",
    "X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y, test_size=0.10, stratify=y, random_state=0)\n",
    "\n",
    "clf = LogisticRegression(random_state=0, solver='lbfgs',multi_class='ovr').fit(X_train, y_train)\n",
    "\n",
    "print(\"Score on test data:\",clf.score(X_test,y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn import metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = clf.predict(X_test)\n",
    "\n",
    "cm = metrics.confusion_matrix(y_test, predictions)\n",
    "print(cm)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
